// --------------------------------------------------------------------------------------------------------------------
// <copyright file="KalmanFilter.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------
namespace Microsoft.Robotics.Numerics
{
    using System.Diagnostics.CodeAnalysis;

    /// <summary>
    /// Class for Kalman Filter
    /// Notation defined below
    /// </summary>
    public class KalmanFilter
    {
        /// <summary>
        /// Delegate for updating new mean based on old mean and control for Kalman EKF
        /// </summary>
        /// <param name="u">The control vector u_{t} of dimension mx1</param>
        /// <param name="meanOld">The mean x_{t-1} of dimension nx1</param>
        /// <returns>The mean x_{t} of dimension nx1</returns>
        [SuppressMessage("StyleCop.CSharp.DocumentationRules", "SA1650:ElementDocumentationMustBeSpelledCorrectly", Justification = "Reviewed. Suppression is OK here.")]
        public delegate double[] PredictMeanDelegate(double[] u, double[] meanOld);

        /// <summary>
        /// Delegate for predicting measurement based on predicted mean for Kalman EKF
        /// </summary>
        /// <param name="meanPred">The predicted mean based on control</param>
        /// <returns>The predicted measurement</returns>
        public delegate double[] PredictMeasurementDelegate(double[] meanPred);

        /// <summary>
        /// Computes one step of the linear Kalman filter 
        /// <para>
        /// Notation follows probablistic robotics section 3.2
        /// Linear Kalman filter is optimal For Linear Gaussian systems defined by following 2 equations 
        ///  1. State transition equation is 
        ///            x_{t} = A_{t} * x_{t-1} + B_{t} * u_{t} + epsilon_{t}
        ///       where 
        ///             x_{t} and x_{t-1} are state vectors at time t and t-1 of dimension nx1
        ///             u_{t} is control vector at time t of dimension mx1
        ///             A_{t} is matrix of dimension nxn, and B_{t} is matrix of dimension nxm.
        ///             epsilon_{t} is a gaussian random vector that models uncertainity introduced by state transition.
        ///             mean of epsilon_{t} is zero, and co-variance is denoted R_{t} where R is of dimension nxn
        ///  2. Measurement equation is
        ///           z_{t} = C_{t} * x_{t} + delta_{t}
        ///       where
        ///           z_{t} is measurement vector of dimension kx1
        ///           C_{t} is matrix of dimension nxk
        ///           delta_{t} is measurement noise
        ///           mean of delta_{t} is zero and variance is denoted Q_{t} where Q is of dimension kxk
        ///  </para>
        ///  <para>
        ///  Linear Kalman Filter operates in two steps
        ///    First, it predicts mean and variance of state based on control alone.
        ///       predMean_{t} = A_{t} * mean_{t-1} + B_{t} * u_{t} 
        ///       predSigma_{t} = A_{t} * sigma_{t-1} * A_{t}' + R_{t}, where ' indicates transpose
        ///    Next, we update mean and variance based on measurement as follows
        ///    First compute Kalmangain K_{t} = predSigma_{t} * C_{t}' * ((C_{t} * predSigma_{t} * C_{t}' + Q_t) ^ (-1))
        ///      The measurement is predicted as follows
        ///        predZ_t = C_{t} * predMean_{t}
        ///     Finally, we update mean and variance of state as follows
        ///       mean_{t} = predMean_{t} + K_{t} * (z_{t} - predZ_t)
        ///        sigma_{t} - (I - K_{t} * C_{t}) * predSigma_{t}
        ///   </para>
        /// </summary>
        /// <param name="meanOld">The mean of state at time t-1</param>
        /// <param name="varianceOld">The variance at time t-1</param>
        /// <param name="u">The control at time t</param>
        /// <param name="z">The measurement at time t</param>
        /// <param name="paramA">The param A at time t</param>
        /// <param name="paramB">The param B at time t</param>
        /// <param name="paramC">The param C at time t</param>
        /// <param name="paramR">The param R at time t</param>
        /// <param name="paramQ">The param Q at time t</param>
        /// <param name="n">The dimension of state</param>
        /// <param name="m">The dimension of control</param>
        /// <param name="k">The dimension of measurement</param>
        /// <param name="meanNew">The mean of state at time t</param>
        /// <param name="varianceNew">The variance of state at time t</param>
        [SuppressMessage("StyleCop.CSharp.DocumentationRules", "SA1650:ElementDocumentationMustBeSpelledCorrectly", Justification = "Reviewed. Suppression is OK here.")]
        public static void ComputeKalmanFilterLinearOneStep(
            double[] meanOld, 
            double[,] varianceOld, 
            double[] u, 
            double[] z, 
            double[,] paramA, 
            double[,] paramB, 
            double[,] paramC, 
            double[,] paramR, 
            double[,] paramQ, 
            int n, 
            int m, 
            int k, 
            out double[] meanNew, 
            out double[,] varianceNew)
        {
            meanNew = null;
            varianceNew = null;

            double[] meanTemp = PredictMeanWithControl(meanOld, paramA, paramB, u);
            double[,] varianceTemp = PredictVarianceWithControl(varianceOld, paramA, paramR);

            double[] measurementPred = PredictMeasurement(meanTemp, paramC);

            UpdateMeanVarianceWithMeasurement(
                paramC, paramQ, n, ref meanNew, ref varianceNew, z, meanTemp, varianceTemp, measurementPred);
        }

        /// <summary>
        /// Computing Extended Kalman Filter (EKF)
        /// <para>
        /// Notation follows probabilistic robotics section 3.3
        /// The extended Kalman Filter (EKF) relaxes the linear assumption 
        ///    and generalizes the linear Gaussian model
        /// The system assumed to be governed by the following two equations
        ///  1. State transition equation is 
        ///               x_{t} = g(u_t, x_{t-1}) + epsilon_{t}
        ///      where
        ///             x_{t} and x_{t-1} are state vectors at time t and t-1 of dimension nx1
        ///             u_{t} is control vector at time t of dimension mx1
        ///             g is a non-linear function of state at previous time, x_{t-1} and the control u_{t}
        ///             epsilon_{t} is a gaussian random vector that models uncertainty introduced by state transition.
        ///             mean of epsilon_{t} is zero, and co-variance is denoted R_{t} where R is of dimension nxn
        ///   2. The measurement equations is 
        ///               z_{t} = h(x_{t}) + delta_{t}
        ///       where
        ///           z_{t} is measurement vector of dimension kx1
        ///           where h is a non-linear function of x_t
        ///           delta_{t} is measurement noise
        ///           mean of delta_{t} is zero and variance is denoted Q_{t} where Q is of dimension kxk
        ///   These equations are linearized via Taylor expansion where
        ///      g(u_{t}, x_{t-1}) ~ g(u_{t}, mean_{t-1}) + G_{t} (x_{t-1} - mean_{t-1})
        ///      where G is the state Jacobian
        ///    and 
        ///      h(x_{t}) ~ h(meanPred_t) + H_{t} (x_{t} - meanPred_{t}))
        ///       where H is the measurement Jacobian
        ///      </para>
        ///      <para>
        ///  Then, the Extended Kalman Filter operates in two steps
        ///    First, it predicts mean and variance of state based on control alone.
        ///       predMean_{t} = g(u_{t}, mean_{t-1})
        ///                  where the PredictMeanDelegate represents g
        ///       predSigma_{t} = G_{t} * sigma_{t-1} * G_{t}' + R_{t}, where ' indicates transpose, and G is state Jacobian
        ///    Next, we update mean and variance based on measurement as follows
        ///    First compute Kalmangain K_{t} = predSigma_{t} * C_{t}' * ((C_{t} * predSigma_{t} * C_{t}' + Q_t) ^ (-1))
        ///      The measurement is predicted as follows
        ///        predZ_t = h(predMean_t) 
        ///                   where the PredictMeasurementDelegate represents h
        ///     Finally, we update mean and variance of state as follows
        ///       mean_{t} = predMean_{t} + K_{t} * (z_{t} - predZ_t)
        ///        sigma_{t} - (I - K_{t} * C_{t}) * predSigma_{t}   
        /// </para>
        /// </summary>
        /// <param name="meanOld">The mean of state at time t-1 </param>
        /// <param name="varianceOld">The sigma of state at time t-1</param>
        /// <param name="u">The control vector u_t</param>
        /// <param name="z">The measurement z_t</param>
        /// <param name="predictMean">The predict Mean delegate</param>
        /// <param name="predictMeasurement">The predict measurement delegate</param>
        /// <param name="stateJacobian">The state Jacobian</param>
        /// <param name="measurementJacobian">The measurement Jacobian</param>
        /// <param name="paramR">The parameter R_t</param>
        /// <param name="paramQ">The parameter Q_t</param>
        /// <param name="n">The dimension of state</param>
        /// <param name="m">The dimension of control</param>
        /// <param name="k">The dimension of measurement</param>
        /// <param name="meanNew">The mean of state at time t</param>
        /// <param name="varianceNew">The variance of state at time t</param>
        [SuppressMessage("StyleCop.CSharp.DocumentationRules", "SA1650:ElementDocumentationMustBeSpelledCorrectly", Justification = "Reviewed. Suppression is OK here.")]
        public static void ComputeKalmanFilterEKFOneStep(
            double[] meanOld, 
            double[,] varianceOld, 
            double[] u, 
            double[] z, 
            PredictMeanDelegate predictMean, 
            PredictMeasurementDelegate predictMeasurement, 
            double[,] stateJacobian, 
            double[,] measurementJacobian, 
            double[,] paramR, 
            double[,] paramQ, 
            int n, 
            int m, 
            int k, 
            out double[] meanNew, 
            out double[,] varianceNew)
        {
            meanNew = null;
            varianceNew = null;

            double[] meanTemp = predictMean(u, meanOld);
            double[,] varianceTemp = PredictVarianceWithControl(varianceOld, stateJacobian, paramR);

            double[] measurementPred = predictMeasurement(meanTemp);

            UpdateMeanVarianceWithMeasurement(
                measurementJacobian, paramQ, n, ref meanNew, ref varianceNew, z, meanTemp, varianceTemp, measurementPred);
        }

        /// <summary>
        /// Predicts variance of state based only on control
        /// </summary>
        /// <param name="varianceOld">The old variance</param>
        /// <param name="paramA">The parameter A in the state change equation x_{t} = A_{t} * x_{t-1} + B_{t} * u_{t} + epsilon_{t}<see cref="T:KalmanFilter"/></param>
        /// <param name="paramR">The variance R of the parameter epsilon_t</param>
        /// <returns>The predicted variance based on control</returns>
        private static double[,] PredictVarianceWithControl(double[,] varianceOld, double[,] paramA, double[,] paramR)
        {
            double[,] varianceTemp;
            double[,] temp3 = paramA.Multiply(varianceOld);

            double[,] temp4 = temp3.Multiply(paramA.Transpose());
            varianceTemp = temp4.Add(paramR);
            return varianceTemp;
        }

        /// <summary>
        /// Predicts mean of state at time t based on control using equation
        /// mean_{t} = A_{t} * mean_{t-1} + B_{t} * u_{t} 
        /// </summary>
        /// <param name="meanOld">The mean of state at time t-1</param>
        /// <param name="paramA">The param A_t</param>
        /// <param name="paramB">The param B_t</param>
        /// <param name="u">The control u_t</param>
        /// <returns>The predicted mean of state based on control</returns>
        private static double[] PredictMeanWithControl(double[] meanOld, double[,] paramA, double[,] paramB, double[] u)
        {
            double[] temp1 = paramA.Multiply(meanOld);
            double[] temp2 = paramB.Multiply(u);

            double[] meanTemp = temp1.Add(temp2);
            return meanTemp;
        }

        /// <summary>
        /// Predicts measurement based on predicted mean at time t
        /// predZ_t = C_{t} * mean_{t}
        /// </summary>
        /// <param name="meanTemp">The predicted mean</param>
        /// <param name="paramC">The parameter C_t</param>
        /// <returns>The predicted measurement</returns>
        private static double[] PredictMeasurement(double[] meanTemp, double[,] paramC)
        {
            double[] measurementPred = paramC.Multiply(meanTemp);
            return measurementPred;
        }

        /// <summary>
        /// Updates mean and variance of state based on measurement as follows
        ///    First compute Kalmangain K_{t} = sigma_{t} * C_{t}' * ((C_{t} * sigma_{t} * C_{t}' + Q_t) ^ (-1))
        ///      The measurement is predicted as follows
        ///        predZ_t = C_{t} * mean_{t}
        ///     Finally, we update mean and variance of state as follows
        ///       mean_{t} = mean_{t} + K_{t} * (z_{t} - predZ_t)
        ///        sigma_{t} - (I - K_{t} * C_{t}) * sigma_{t}
        /// </summary>
        /// <param name="paramC">The parameter C</param>
        /// <param name="paramQ">The parameter Q</param>
        /// <param name="n">The dimension of state </param>
        /// <param name="meanNew">The updated mean based on measurement</param>
        /// <param name="varianceNew">The updated variance based on control</param>
        /// <param name="z">The measurement z</param>
        /// <param name="meanTemp">The predicted mean based on control</param>
        /// <param name="varianceTemp">The predicted variance based on control</param>
        /// <param name="measurementPred">The predicted measurement</param>
        private static void UpdateMeanVarianceWithMeasurement(
            double[,] paramC, 
            double[,] paramQ, 
            int n, 
            ref double[] meanNew, 
            ref double[,] varianceNew, 
            double[] z, 
            double[] meanTemp, 
            double[,] varianceTemp, 
            double[] measurementPred)
        {
            double[,] kalmanGain = ComputeKalmanGain(paramC, paramQ, varianceTemp);

            // Computing Kalman gain
            if (kalmanGain == null)
            {
                return;
            }

            double[] innovation = z.Subtract(measurementPred);

            double[] scaledInnovation = kalmanGain.Multiply(innovation);
            meanNew = meanTemp.Add(scaledInnovation);

            double[,] identityMatrix = MatrixExtensions.GetIdentityMatrix(n);
            double[,] temp7 = kalmanGain.Multiply(paramC);
            double[,] scaleFactor = identityMatrix.Subtract(temp7);
            varianceNew = scaleFactor.Multiply(varianceTemp);
        }

        /// <summary>
        /// Computes Kalman gain 
        /// </summary>
        /// <param name="paramC">The parameter C_t of dimension kxn which maps state to measurement</param>
        /// <param name="paramQ">The observataion variance Q_t which is kxk matrix</param>
        /// <param name="varianceTemp">The state variance_t computed based only on control</param>
        /// <returns>The Kalman gain</returns>
        [SuppressMessage("StyleCop.CSharp.DocumentationRules", "SA1650:ElementDocumentationMustBeSpelledCorrectly", Justification = "Reviewed. Suppression is OK here.")]
        private static double[,] ComputeKalmanGain(double[,] paramC, double[,] paramQ, double[,] varianceTemp)
        {
            double[,] kalmanGain = null;
            double[,] sigmaTimesCTran = varianceTemp.Multiply(paramC.Transpose());

            double[,] temp5 = MatrixExtensions.Multiply(paramC, sigmaTimesCTran);
            double[,] innovationCovariance = MatrixExtensions.Add(temp5, paramQ);

            double[,] inverse = innovationCovariance.Invert();

            if (inverse != null)
            {
                kalmanGain = sigmaTimesCTran.Multiply(inverse);
            }

            return kalmanGain;
        }
    }
}
